from langchain_core.tools import create_retriever_tool

from documents.dense_insert_milvus_optimize import MilvusVectorSave

mv = MilvusVectorSave()
mv.create_connection()
retriever = mv.vector_store_saved.as_retriever(
    search_type='similarity',  # 仅返回相似度超过阈值的文档
    search_kwargs = {
        "k": 4,
        "score_threshold": 0.1,
        "ranker_type": "rrf",
        "ranker_params": {"k": 100},
        # 'filter': {"category": "content"}
    }
)

retriever_tool = create_retriever_tool(
    retriever=retriever,
    name='rag_retriever',
    description='搜索并返回关于 ‘半导体和芯片’ 的信息, 内容涵盖：半导体和芯片的封装、测试、光刻胶等'
)